import pandas as pd
import joblib

# 加载模型
model = joblib.load('../model/best_xgb_pipeline_model.joblib')

# 新样本
new_sample = pd.DataFrame({
    'Hormonal Contraceptives (years)': [0],
    'Dx:Cancer': [0],
    'STDs (number)': [0],
    'STDs:genital herpes': [1],
    'STDs:HIV': [0],
    'STDs:condylomatosis': [0],
    'STDs': [0],
    'Hinselmann': [0],
    'Dx:CIN': [0],
    'Citology': [1],
    'IUD': [0],
    'IUD (years)': [0],
    'STDs:AIDS': [0],
    'Number of sexual partners': [2],
    'Age': [20],
    'STDs:vaginal condylomatosis': [0],
    'STDs:vulvo-perineal condylomatosis': [0],
    'Smokes (years)': [0],
    'First sexual intercourse': [22],
    'Schiller': [0]
}, columns=[
    'Hormonal Contraceptives (years)', 'Dx:Cancer', 'STDs (number)', 'STDs:genital herpes',
    'STDs:HIV', 'STDs:condylomatosis', 'STDs', 'Hinselmann', 'Dx:CIN', 'Citology',
    'IUD', 'IUD (years)', 'STDs:AIDS', 'Number of sexual partners', 'Age',
    'STDs:vaginal condylomatosis', 'STDs:vulvo-perineal condylomatosis',
    'Smokes (years)', 'First sexual intercourse', 'Schiller'
])

# 预测类别
pred_label = model.predict(new_sample)[0]

# 预测所有概率
pred_proba = model.predict_proba(new_sample)[0]

# 输出判为该类别时的概率
print(f"新样本预测类别：{pred_label}")
print(f"新样本属于类别 {pred_label} 的概率：{pred_proba[pred_label]}")